import torch
from torchvision import transforms as T
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from PIL import Image
import matplotlib.pyplot as plt


class ImageDs(Dataset):
    def __init__(self):
        super().__init__()
        self.df = pd.read_csv('meta.csv', encoding="gbk")
        self.transformer = T.Compose([
            T.Resize((100, 100), antialias=True),
            # T.Grayscale(),
            T.ToTensor()
        ])
        self.label_map = ['Apple', 'Banana', 'Orange']

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # 根据idx获取一行数据，再拆分出路径 和 标签
        row = self.df.iloc[idx]  # 注意这里是[] 不是小括号
        img_path = row["img_path"]
        label = row["label"]
        # 路径用于PIL读取图像然后进行transformer转换为tensor张量
        img = Image.open(img_path).convert("RGB")
        if self.transformer:
            img = self.transformer(img)
        # 标签用于label_map映射为 0 和 1、2的张量
        label = torch.tensor(self.label_map.index(label))
        return img, label


# 实例化对象
ds = ImageDs()

# 打印长度
print(len(ds))

# 找到图像和标签看一下
# img, label = ds[0]
# print(img.shape, label)

# 随机批量加载
dl = DataLoader(ds, batch_size=4, shuffle=True)
print(next(enumerate(dl)))

# 批量可视化
for i, (img, label) in enumerate(dl):
    plt.figure(figsize=(10, 10))
    for j in range(img.shape[0]):
        plt.subplot(2, 2, j + 1)
        plt.imshow(img[j].permute(1, 2, 0))
        plt.title(ds.label_map[label[j]])
        plt.axis("off")
plt.show()
